import numpy as np
import paddlehub as hub
import cv2
from matplotlib import image as mpimg, pyplot as plt

from src.config.config import TEST_PATH

module = hub.Module(name="resnet50_vd_animals")

# 以只读方式打开test.txt文件
with open(TEST_PATH, 'r', encoding='utf-8') as f:
    test_img_path = []  # 创建待预测图片路径空列表
    for line in f:  # 遍历test.txt文件
        test_img_path.append(line.strip())  # strip函数去除图片路径首尾无用的空格，并依次添加到列表末尾


# 以二进制流读取文件（支持中文路径）
def read_image_utf8(image_path):
    raw_data = np.fromfile(image_path, dtype=np.uint8)
    # 解码图像数据
    img = cv2.imdecode(raw_data, cv2.IMREAD_COLOR)  # 使用cv2.IMREAD_UNCHANGED保留Alpha通道
    if img is None:
        raise ValueError(image_path + "图像解码失败，请检查文件路径或格式")
    return img


# 展示其中大熊猫图片
# img1 = mpimg.imread(test_img_path[0])
# plt.figure(figsize=(10, 10))
# plt.imshow(img1)
# plt.axis('off')
# plt.show()

# 从路径列表中依次读取待预测图片
np_images = [read_image_utf8(image_path) for image_path in test_img_path]

# 调用module api实现预训练模型预测图片中的动物
results = module.classification(images=np_images)

# 打印预测结果
for result in results:
    print(result)
    # print(list(result.keys())[0])
